"""
这是一个模拟空战多对多的环境，参照gym的交互逻辑设计，敌方目标随机生成于半径为150km的球内(已知歼11机载雷达的最大探测距离)，分辨率为1公里
"""

import math
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import axes3d
import Env.Tools.ACFunc as ACF
from Env.SAM import SAM
from Env.UAV import UAV
import Env.Tools.ExcelFunc as exl

Unit = {'UAV': UAV,
        'SAM': SAM}


class ACEnv:
    def __init__(self, n_friend, n_enemy,
                 Battle_range=40000, height=20000, 初始生成距离=30000,
                 最大时间=3000, type_e='SAM', type_f='UAV', 时间步长=0.1):
        self.action_space = range(72)  # 此处为所有初始步可能的动作
        self.Battle_range = Battle_range  # 战场平面范围 m
        self.height = height  # 高度限制 m
        # self.n_target=n_target  # 目标数目，暂时认为敌方既为目标
        self.type_e=type_e
        self.n_enemy = n_enemy
        self.type_f=type_f
        self.n_friend = n_friend
        self.duration = 最大时间  # 单位秒
        self.dt = 时间步长
        self.rou = math.ceil(最大时间 / 时间步长)
        self.init_distance = 初始生成距离
        self.exc = exl.excel()  # 构建记录文件


        """
        self.att_target = pd.DataFrame(np.zeros([n_friend,n_enemy],dtype=int),
                                     index=['第{}架友机'.format(str(i)) for i in range(1,(n_friend+1))],
                                     columns=['第{}个目标'.format(str(i)) for i in range(1,(n_enemy+1))])
        
        坐标位置放在每个实体中
        self.Location_F=pd.DataFrame(np.zeros([n_friend,3]),
                                     index=['第{}架友机'.format(str(i)) for i in range(1,(n_friend+1))],
                                     columns=list('XYZ'))   # 友机方位矩阵
        self.Location_E = pd.DataFrame(np.zeros([n_enemy, 3]),
                                       index=['第{}架敌机'.format(str(i)) for i in range(1, (n_enemy + 1))],
                                       columns=list('XYZ'))  # 敌机方位矩阵
        """

    def reset(self):    # TODO （zrc） 场景初始化，随机敌我双方出生点，返回
        self.enemies = self.objs_creat(self.n_enemy, self.type_e, 1)  # 单位类型数组
        self.friends = self.objs_creat(self.n_friend, self.type_f, 0)  # 单位类型数组
        self.att_target_f = np.array([[-1] * self.n_friend] * self.n_enemy)  # 导弹命中倒计时f
        self.att_target_e = np.array([[-1] * self.n_enemy] * self.n_friend)  # 导弹命中倒计时e
        self.att_result = [[0] * self.n_enemy, [0] *self.n_friend]  # 攻击结果记录
        self.exc.creat_sheet()
        self.build()  # 依据参数构建场景

        observation=self.get_observation()
        return observation

    def FirstStep(self, schedule):  # TODO (zrc) 独立的第一步动作，确定在出生点时怎么进行*目标分配*，分配后直接拉到攻击边界
        Points_near = self.离散球面点穷举(20)  # 内球面，半径20
        Points_far = self.离散球面点穷举(22)  # 外球面，半径22

        for i in range(len(self.friends)):
            if schedule[i][1] < 60:
                pn=Points_near[schedule[i][1]]
                pe=self.enemies[schedule[i][0]].Position
                pf=pn+pe
                self.friends[i].Position = pf
                self.friends[i].pointing=[math.pi/2,math.atan((pe[1]-pf[1])/(pe[0]-pf[0]))]     # [俯仰角，偏转角]
            else:
                pn=Points_far[schedule[i][1] - 60]
                pe=self.enemies[schedule[i][0]].Position
                pf=pn+pe
                self.friends[i].Position = pf
                self.friends[i].pointing = [math.pi/2,math.atan((pe[1]-pf[1])/(pe[0]-pf[0]))]

        observation=self.get_observation()
        return observation


    def step(self, action):     # TODO (zrc) 第二阶段的动作，主要接受机动动作（俯仰角、偏转角）和攻击动作
        self.analy_location(action[0])  # 分析机动动作的影响
        self.analy_attack(action[1],0)    # 分析攻击动作的影响,0为我方,1为AI
        observation_, done = self.do()  # 进行动作仿真，得到observation_和done
        reward = self.reward()  # 计算回报值
        if done:
            self.exc.save()
        return observation_, reward, done

    def render(self):  # TODO (zrc) 展示当前敌我双方位置
        Location_F = [f.Position for f in self.friends]  # 己方坐标
        Location_E = [e.Position for e in self.enemies]  # 敌方坐标
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        x_f = [i[0] for i in Location_F]
        y_f = [i[1] for i in Location_F]
        z_f = [i[2] for i in Location_F]
        x_e = [i[0] for i in Location_E]
        y_e = [i[1] for i in Location_E]
        z_e = [i[2] for i in Location_E]
        ax.scatter3D(x_f, y_f, z_f, c='r', marker='*')  # TODO (zrc) 我方是红色，对方蓝色
        ax.scatter3D(x_e, y_e, z_e, c='b', marker='D')
        ax.set_zlim3d(0, self.height/1000)
        plt.show()

    def build(self, View=1):  # 随机生成敌我双方位置,View为视角，1为以敌方为中心，0为以我方为中心
        mu1 = self.球面点随机(self.init_distance) if View == 1 else [0., 0., 0.]
        mu2 = [0., 0., 0.] if View == 1 else self.球面点随机(self.init_distance)
        for i in self.friends:
            i.Position = ((mu1 + np.random.uniform(-500, 500, [1, 3]))/1000)[0]

        for j in self.enemies:
            ran = np.random.uniform(-5000, 5000, [1, 3])
            ran[0, 2] = 0.
            j.Position = ((mu2 + ran)/1000)[0]

    def analy_location(self,flight):  # 分析动作会对环境哪些方面产生影响，进行拆分
        a=0
        for i in self.friends:
            i.FlightTo(flight[a][0],flight[a][1],self.dt)
            a+=1

    def analy_attack(self,attack,D_W):  # TODO　（ｚｒｃ）　接收到的参数为[[0，1，0，1]...[...]]4*4的矩阵
        a=-1
        for att in attack:
            a+=1
            if D_W == 0:    # 我方
                for i in range(len(att)):
                    if att[i]==1 and self.friends[a].Attack(self.enemies[i].Position):
                        self.att_target_f[a][i]=\
                            ((self.distance_xyz(self.friends[a].Position,
                                                self.enemies[i].Position)/self.dt)//
                             self.friends[a].v_missile)+1

            else:   # 敌方
                for i in range(len(att)):
                    if att[i]==1 and self.enemies[a].Attack(self.friends[i].Position,self.dt):
                        self.att_target_e[a][i]=\
                            (self.distance_xyz(
                            self.enemies[i].Position,self.friends[a].Position)/self.dt)//\
                            self.enemies[a].v_missile+1


    def do(self):  # 对环境进行更改，返回observation_和done
        done=False
        if not ((self.att_target_f == np.array([[-1] * self.n_friend] * self.n_enemy)).all()
                        and (self.att_target_e == np.array([[-1] * self.n_enemy] * self.n_friend)).all()):
            self.att_result=self.attack_result()     # TODO （zrc）根据结果处理成我方地方存活情况
            destroy_e=np.where(self.att_result[0]==2)    # 敌方的被摧毁情况
            destroy_f=np.where(self.att_result[1]==2)    # 我方的被摧毁情况
            for i in destroy_f[1]:
                self.friends[i].live=0
            for j in destroy_e[1]:
                self.enemies[j].live=0

            self.flash_att_target()

        livelist = [[f.live for f in self.friends], [e.live for e in self.enemies]]

        if (np.array(livelist[0])==0).all() or (np.array(livelist[1])==0).all():
            done=True

        observation=self.get_observation()
        return observation,done

    def reward(self):  # 回报计算，是否命中等
        pass

    def objs_creat(self, num, type, D_W):  # 单位创建,D_W为敌我识别,0是我方，1是对方
        obj = []
        for i in range(0, num):
            name = 'FF' + str(i + 1) if D_W == 0 else 'DF'
            temp = Unit[type](name)
            obj.append(temp)
        return obj


    def attack_result(self):  # TODO （zrc）输入我方攻击列表和对方攻击列表，输出命中，未命中和差点命中
        # TODO （zrc）刷新状态,计算武器概率

        List_F = [f.Position for f in self.friends]
        for i in range(len(List_F), 4):
            List_F.append([0, 0, 0])
        Location_F = np.array(List_F)  # TODO (zrc) 将坐标补齐成4个,下同

        List_E = [e.Position for e in self.enemies]
        for i in range(len(List_E), 4):
            List_E.append([0, 0, 0])
        Location_E = np.array(List_E)

        Fire_para_F = ACF.utility(4, 4, Location_F, Location_E)
        Atta_Matrix_F = Fire_para_F.Run()  # 返回攻击概率矩阵
        Fire_para_E = ACF.utility(4, 4, Location_E, Location_F)
        Atta_Matrix_E = Fire_para_E.Run()  # 返回攻击概率矩阵
        Att_Result_F = ACF.Attack.att_result(Atta_Matrix_F, self.att_target_f)
        Att_Result_E = ACF.Attack.att_result(Atta_Matrix_E, self.att_target_e)
        return [Att_Result_F, Att_Result_E]     # 返回攻击结果

    def 球面点随机(self, distance):  # TODO (zrc) 在半球面上随机生成一个点，再转换成直角坐标系
        Theta = np.random.uniform(0, math.pi * 0.5)
        Phi = np.random.uniform(0, math.pi * 2)
        x = distance * math.sin(Theta) * math.cos(Phi)
        y = distance * math.sin(Theta) * math.sin(Phi)
        z = distance * math.cos(Theta)
        return [x, y, z]

    def 离散球面点穷举(self, r):
        points = []
        for i in range(1, 6):  # theta取值 15,30,45,60,75
            for j in range(12):  # phi取值0-360，步长30
                Theta = math.pi / 12 * i
                Phi = math.pi / 6 * j
                x = r * math.sin(Theta) * math.cos(Phi)
                y = r * math.sin(Theta) * math.sin(Phi)
                z = r * math.cos(Theta)
                points.append([x, y, z])
        return points

    def 展示(self):
        fig=plt.figure()
        plot=fig.add_subplot(111,projection='3d')
        X_F = [i.Position[0][0] for i in self.friends]
        Y_F = [i.Position[0][1] for i in self.friends]
        Z_F = [i.Position[0][2] for i in self.friends]
        X_E = [i.Position[0][0] for i in self.enemies]
        Y_E = [i.Position[0][1] for i in self.enemies]
        Z_E = [i.Position[0][2] for i in self.enemies]
        plot.scatter3D(X_F, Y_F, Z_F, c='r', marker='*')
        plot.scatter3D(X_E, Y_E, Z_E, c='b', marker='8')

        plt.show()

    def distance(self,x1,x2,y1,y2,z1=0,z2=0):
        dis=math.sqrt((x2-x1)**2+(y2-y1)**2+(z2-z1)**2)
        return dis
    def distance_xyz(self,A,B):
        x=B[0]-A[0]
        y = B[1] - A[1]
        z = B[2] - A[2]
        dis=math.sqrt(x**2+y**2+z**2)
        return dis

    def flash_att_target(self):     # 更新att_target列表
        self.att_target_f-=1
        self.att_target_e-=1
        f_list=np.where(self.att_target_f==-2)
        e_list=np.where(self.att_target_e==-2)
        for i,j in zip(f_list[0],f_list[1]):
            self.att_target_f[i][j]=-1
        for i,j in zip(e_list[0],e_list[1]):
            self.att_target_e[i][j]=-1

    def get_observation(self):  # 得到observation
        observation=[]
        Location_F = [f.Position for f in self.friends]  # 己方坐标
        Location_E = [e.Position for e in self.enemies]  # 敌方坐标
        n_ammo = [f.NumberOfMissile for f in self.friends]  # 载弹量
        livelist = [[f.live for f in self.friends], [e.live for e in self.enemies]]

        if (np.array(livelist[0]) == 0).all() or (np.array(livelist[1]) == 0).all():
            done = True

        observation.append(Location_F)  # 我方位置
        observation.append(Location_E)  # 敌方位置
        observation.append(n_ammo)  # 我方弹药情况
        observation.append(self.att_target_f)  # 攻击状态
        observation.append(self.att_target_e)  # 攻击状态
        observation.append(livelist)  # 场上生存情况
        observation.append(self.att_result[0])  # 我方i对敌方j的摧毁情况
        observation.append(self.att_result[1])  # 敌方i对我方j的摧毁情况

        self.exc.out_to_excel(observation)
        return observation

    def excel_save(self):
        self.exc.save()



